// prefetcher.c (修正后)
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <time.h>
#include <stdint.h>
#include <sys/resource.h>
#include <bpf/libbpf.h>
#include <linux/limits.h> // For PATH_MAX

#include "prefetcher.h"
#include "prefetcher.skel.h"

// ===================================
//      仅在用户态使用的定义 (增加新状态和宏)
// ===================================
#define HISTORY_SIZE 8
#define DETECTION_THRESHOLD 3
#define TIME_WINDOW_NS (50 * 1000 * 1000)
#define PAGE_SIZE 4096

// *** 优化点1: 更激进的预取策略 ***
// 当模式确认后，一次性预取未来的64个步长
#define BATCH_PREFETCH_COUNT 64

typedef enum {
    DETECTING,  // 探测模式中
    CONFIRMED,  // 模式刚被确认，准备执行一次性批量预取
    PREFETCHING // 已经执行了批量预取，现在处于观察/冷却状态
} pattern_state_t;

typedef struct {
    pattern_state_t state;
    int stable_count;
    int64_t current_stride;

    uint64_t addr_history[HISTORY_SIZE];
    uint64_t ts_history[HISTORY_SIZE];
    int history_idx;

    // *** 优化点2: 记录预取范围 ***
    uint64_t prefetch_end_addr; // 记录上次批量预取的结束地址

    time_t last_event_time;
} process_state_t;


// 全局变量
static volatile bool exiting = false;
static process_state_t *state = NULL;
static char target_file[PATH_MAX] = {0};

// 信号处理函数
static void sig_handler(int sig) {
    exiting = true;
}

// 初始化进程状态机
process_state_t* init_process_state() {
    process_state_t *s = (process_state_t*)calloc(1, sizeof(process_state_t));
    if (!s) {
        perror("Failed to allocate process state");
        return NULL;
    }
    s->state = DETECTING; // 初始状态是探测
    s->prefetch_end_addr = 0;
    s->last_event_time = time(NULL);
    return s;
}

// 虚拟地址到文件偏移的转换
static int vma_to_file_offset(pid_t pid, uint64_t vma, char *filepath, size_t filepath_len, off_t *offset) {
    char maps_path[64];
    snprintf(maps_path, sizeof(maps_path), "/proc/%d/maps", pid);

    FILE *f = fopen(maps_path, "r");
    if (!f) return -1;

    char *line = NULL;
    size_t len = 0;
    ssize_t read;
    int found = -1;

    while ((read = getline(&line, &len, f)) != -1) {
        uint64_t start, end, file_off;
        char perms[5], dev[6];
        long inode;
        char current_path[PATH_MAX] = {0};
        
        int items = sscanf(line, "%lx-%lx %4s %lx %5s %ld %s", &start, &end, perms, &file_off, dev, &inode, current_path);
        
        if (vma >= start && vma < end && inode != 0 && strchr(perms, 'r') && items == 7) {
            strncpy(filepath, current_path, filepath_len - 1);
            filepath[filepath_len - 1] = '\0';
            *offset = file_off + (vma - start);
            found = 0;
            break;
        }
    }

    free(line);
    fclose(f);
    return found;
}

// 核心处理逻辑 (*** 完全重写 ***)
void process_event(const struct page_fault_event *e) {
    char current_filepath[PATH_MAX];
    off_t current_offset;

    if (vma_to_file_offset(e->pid, e->addr, current_filepath, sizeof(current_filepath), &current_offset) != 0) {
        return;
    }
    if (strcmp(target_file, current_filepath) != 0) {
        return;
    }

    printf("Event: PID=%u, Addr=0x%llx, State: %d\n", e->pid, e->addr, state->state);

    // 状态机核心逻辑
    switch (state->state) {
        case DETECTING: {
            int history_idx = state->history_idx;
            int prev_idx = (history_idx - 1 + HISTORY_SIZE) % HISTORY_SIZE;
            uint64_t last_addr = state->addr_history[history_idx == 0 ? 0 : prev_idx];
            uint64_t last_ts = state->ts_history[history_idx == 0 ? 0 : prev_idx];

            state->addr_history[history_idx] = e->addr;
            state->ts_history[history_idx] = e->timestamp;
            state->history_idx = (history_idx + 1) % HISTORY_SIZE;

            if (history_idx < 1) return;

            if ((e->timestamp - last_ts) > TIME_WINDOW_NS) {
                state->stable_count = 0;
                state->current_stride = 0;
                return;
            }

            int64_t stride = e->addr - last_addr;
            if (stride == 0) return;

            printf("  -> Detecting... Stride: %ld\n", stride);

            if (stride == state->current_stride) {
                state->stable_count++;
            } else {
                state->current_stride = stride;
                state->stable_count = 1;
            }

            if (state->stable_count >= DETECTION_THRESHOLD) {
                printf("  ==> Pattern CONFIRMED! Stride: %ld. Switching to CONFIRMED state.\n", state->current_stride);
                state->state = CONFIRMED; // 切换到待预取状态
            }
            break;
        }

        case CONFIRMED: {
            printf("  ==> Executing SYNC BATCH PREFETCH for event at 0x%lx.\n", e->addr);
            
            // *** 核心修正：在阻塞前，先更新历史记录 ***
            // 这确保了下一次事件到来时，'last_addr' 是正确的
            state->addr_history[state->history_idx] = e->addr;
            state->history_idx = (state->history_idx + 1) % HISTORY_SIZE;

            // *** 然后，立即切换到 PREFETCHING 状态 ***
            // 这样，在 read 循环期间如果收到新事件（虽然我们处理不了），它们至少知道我们正在预取
            state->state = PREFETCHING;

            int fd = open(target_file, O_RDONLY);
            if (fd < 0) {
                fprintf(stderr, "WARN: Could not open target file %s for prefetching: %s\n", target_file, strerror(errno));
                state->state = DETECTING; // 出错了，退回探测状态
                return;
            }

            char dummy_buffer[1];
            size_t prefetched_size = 0;
            printf("  ==> Prefetching by reading 1 byte from each predicted page...\n");

            for (int i = 0; i < BATCH_PREFETCH_COUNT; i++) {
                // 从当前事件地址 e->addr 开始预测，而不是 current_offset
                off_t prefetch_offset = current_offset + (off_t)(i + 1) * state->current_stride;
                if (pread(fd, dummy_buffer, 1, prefetch_offset) <= 0) {
                    break;
                }
                prefetched_size += abs(state->current_stride);
            }
            
            close(fd);
            
            state->prefetch_end_addr = e->addr + prefetched_size;
            printf("  ==> PREFETCHING finished. Prefetched until address ~0x%lx\n", state->prefetch_end_addr);
            
            break;
        }

        case PREFETCHING: {
            // 在这个状态下，我们收到的缺页事件都是“不应该发生”的
            printf("  -> WARNING: Page fault occurred at 0x%lx during PREFETCHING state.\n", e->addr);
            
            // 简单处理：直接退回探测状态，因为任何缺页都意味着预测失败或模式改变
            state->state = DETECTING;
            
            // 用当前这个意外的事件，作为新一轮探测的开始
            memset(state->addr_history, 0, sizeof(state->addr_history));
            memset(state->ts_history, 0, sizeof(state->ts_history));
            state->history_idx = 0;
            state->stable_count = 0;
            
            state->addr_history[state->history_idx] = e->addr;
            state->ts_history[state->history_idx] = e->timestamp;
            state->history_idx++;
            
            break;
        }
    }
}


// Ring Buffer 回调
static int handle_event_wrapper(void *ctx, void *data, size_t data_sz) {
    const struct page_fault_event *e = data;
    process_event(e);
    return 0;
}

// main 函数
int main(int argc, char **argv) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <pid> <target_file_path>\n", argv[0]);
        fprintf(stderr, "Example: %s 12345 /path/to/testfile.dat\n", argv[0]);
        return 1;
    }
    pid_t pid = atoi(argv[1]);
    if (realpath(argv[2], target_file) == NULL) {
        fprintf(stderr, "Error resolving path for %s: %s\n", argv[2], strerror(errno));
        return 1;
    }

    struct prefetcher_bpf *skel;
    struct ring_buffer *rb = NULL;
    int err;

    libbpf_set_strict_mode(LIBBPF_STRICT_ALL);

    struct rlimit r = {RLIM_INFINITY, RLIM_INFINITY};
    if (setrlimit(RLIMIT_MEMLOCK, &r)) {
        perror("setrlimit(RLIMIT_MEMLOCK)");
        return 1;
    }

    skel = prefetcher_bpf__open();
    if (!skel) {
        fprintf(stderr, "Failed to open BPF skeleton\n");
        return 1;
    }

    skel->rodata->target_pid = pid;

    err = prefetcher_bpf__load(skel);
    if (err) {
        fprintf(stderr, "Failed to load and verify BPF skeleton\n");
        goto cleanup;
    }

    err = prefetcher_bpf__attach(skel);
    if (err) {
        fprintf(stderr, "Failed to attach BPF skeleton\n");
        goto cleanup;
    }

    printf("eBPF prefetcher started. Watching PID: %d, File: %s. Press Ctrl+C to exit.\n", pid, target_file);

    state = init_process_state();
    if (!state) goto cleanup;

    rb = ring_buffer__new(bpf_map__fd(skel->maps.events), handle_event_wrapper, NULL, NULL);
    if (!rb) {
        err = -1;
        fprintf(stderr, "Failed to create ring buffer\n");
        goto cleanup;
    }
    
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);

    while (!exiting) {
        err = ring_buffer__poll(rb, 100 /* 100ms timeout */);
        if (err == -EINTR) {
            err = 0;
            break;
        }
        if (err < 0) {
            printf("Error polling ring buffer: %d\n", err);
            break;
        }
    }

cleanup:
    printf("\nDetaching and cleaning up...\n");
    ring_buffer__free(rb);
    prefetcher_bpf__destroy(skel);
    free(state);
    return -err;
}